#include "PlaneMesh.cuh"

using namespace OBR;

/**
 * @brief do intersection test on plane
 *
 * @param ro ray origin
 * @param rd ray direction
 * @param plane plane data
 * @param u output texcoord[0]
 * @param v output texcoord[1]
 * @return __forceinline__
 */
__forceinline__ __device__ bool intersectPlane(float3 ro, float3 rd, PlaneMeshData &plane, float *t,
                                               float *u, float *v)
{
    float NdotD = dot(plane.n, rd);
    float NdotR = dot(plane.n, plane.o - ro);

    if (NdotD * NdotR > 0)
    {
        *t = NdotR / NdotD;
        float3 dp = ro + rd * (*t) - plane.o;

        // Compute $(u,v)$ offsets at auxiliary points
        // Choose two dimensions to use for ray offset computation
        int dim[2];
        if (fabs(plane.n.x) > fabs(plane.n.y) && fabs(plane.n.x) > fabs(plane.n.z))
        {
            dim[0] = 1;
            dim[1] = 2;
        }
        else if (fabs(plane.n.y) > fabs(plane.n.z))
        {
            dim[0] = 0;
            dim[1] = 2;
        }
        else
        {
            dim[0] = 0;
            dim[1] = 1;
        }

        // Initialize _A_, _B_ matrices for offset computation
        float A[2][2] = {get_float3_at(&(plane.u), dim[0]), get_float3_at(&(plane.v), dim[0]),
                         get_float3_at(&(plane.u), dim[1]), get_float3_at(&(plane.v), dim[1])};
        float B[2] = {get_float3_at(&dp, dim[0]), get_float3_at(&dp, dim[1])};
        if (!solve_linear_system_2x2(A, B, u, v))
            *u = *v = 0;

        if (*u > 0 && *u < 1 && *v > 0 && *v < 1)
        {
            return true;
        }
    }
    return false;
}

extern "C" __global__ void __intersection__plane_mesh()
{
    HitData *sbt_record = reinterpret_cast<HitData *>(optixGetSbtDataPointer());
    PlaneMeshData &data =
        reinterpret_cast<PlaneMeshData *>(sbt_record->hit_payload)[sbt_record->mesh_instance_id];

    const float3 orig = optixGetObjectRayOrigin();
    const float3 dir = optixGetObjectRayDirection();

    float u, v, t;
    if (intersectPlane(orig, dir, data, &t, &u, &v))
    {
        optixReportIntersection(t, 0, float_as_uint(u), float_as_uint(v));
    }
}

extern "C" __global__ void __closesthit__plane_mesh()
{
    uint p0, p1;
    SamplePRD *prd = getPRDWithPayload(p0, p1);
    prd->hit = true;

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("==== plane hit program launched\n");
    }
#endif

    // SBT info
    HitData *sbt_record = reinterpret_cast<HitData *>(optixGetSbtDataPointer());
    const uint mesh_id = sbt_record->mesh_id;
    const uint mesh_instance_id = sbt_record->mesh_instance_id;

    // trace info
    float3 ray_dir = optixGetWorldRayDirection();
    float3 ray_orig = optixGetWorldRayOrigin();
    float ray_t = optixGetRayTmax();
    float3 hit_pos = ray_orig + ray_t * ray_dir;

    // init surface interaction
    InteractionData &si = prd->inter;
    si.type = InteractionType::SURFACE_INTERACTION;
    si.wo = prd->inter.wi;
    si.wo.t = ray_t;
    si.mat_id = sbt_record->mat_id;
    si.mat_instance_id = sbt_record->mat_instance_id;
    si.light_id = sbt_record->light_id;
    si.light_instance_id = sbt_record->light_instance_id;
    si.med_id_inside = sbt_record->med_id_inside;
    si.med_instance_id_inside = sbt_record->med_instance_id_inside;
    si.med_id_outside = sbt_record->med_id_outside;
    si.med_instance_id_outside = sbt_record->med_instance_id_outside;
    si.mesh_id = mesh_id;
    si.mesh_instance_id = mesh_instance_id;

    /* ------------------- calculate surface interaction -------------------- */
    PlaneMeshData &meshdata =
        reinterpret_cast<PlaneMeshData *>(sbt_record->hit_payload)[mesh_instance_id];
    si.n = normalize(optixTransformNormalFromObjectToWorldSpace(meshdata.n));
    // si.n = normalize(world2object.transpose() * meshdata.n);
    float u = uint_as_float(optixGetAttribute_0());
    float v = uint_as_float(optixGetAttribute_1());
    si.uv = make_float2(u, v);
    si.world_pos = hit_pos;
    si.obj_pos = optixTransformPointFromWorldToObjectSpace(hit_pos);

    if (prd->ray_occlusion_test)
    {
        // in occlusion test we dont need differentials
        return;
    }

    /* ----------------------- calculate differentials ---------------------- */
    if (si.wo.has_differential)
    {
        si.dndu = make_float3(0, 0, 0);
        si.dndv = make_float3(0, 0, 0);
        si.dpdu = meshdata.u;
        si.dpdv = meshdata.v;
        float du, dv, dt;
        if (intersectPlane(si.wo.rx_o, si.wo.rx_d, meshdata, &dt, &du, &dv))
        {
            si.dpdx = si.wo.rx_o + si.wo.rx_d * dt - si.world_pos;
            si.dudx = du - u;
            si.dvdx = dv - u;
        }
        else
        {
            si.dpdx = make_float3(0, 0, 0);
            si.dudx = si.dvdx = 0;
        }
        if (intersectPlane(si.wo.ry_o, si.wo.ry_d, meshdata, &dt, &du, &dv))
        {
            si.dpdy = si.wo.ry_o + si.wo.ry_d * dt - si.world_pos;
            si.dudy = du - u;
            si.dvdy = dv - u;
        }
        else
        {
            si.dpdx = make_float3(0, 0, 0);
            si.dudx = si.dvdx = 0;
        }
    }

    return;
}